"""
模型评估
"""
import numpy as np
import torch
from torch.optim import Adam
from tqdm import tqdm

from chatbot import config
from chatbot.dnn.sort.dataset import dnnsort_data_loader
from chatbot.dnn.sort.siamese import SiameseNetwork

import torch.nn.functional as F


def eval(by_word=True):
    model = SiameseNetwork().to(config.device)
    model.load_state_dict(torch.load(config.sort_model_save_path))
    optimizer = Adam(model.parameters(), lr=0.001)
    optimizer.load_state_dict(torch.load(config.sort_optimizer_save_path))

    bar = tqdm(enumerate(dnnsort_data_loader), total=len(dnnsort_data_loader))
    loss_list=[]
    acc_list=[]
    for idx, (input1, input2, target) in bar:
        output = model(input1, input2)
        loss = F.nll_loss(output, target)
        loss_list.append(loss.item())
        pred=torch.max(output,dim=-1)[-1]
        acc=pred.eq(target).float().mean()
        acc_list.append(acc)
    print(np.mean(loss_list),np.mean(acc_list))

